{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Introduction\n",
    "\n",
    "This is a simple tool (beta) to convert your own datasets to ConKADI-format datasets.\n",
    "\n",
    "This tool includes two notebooks; this is the first one, which generates middle-format datasets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "ename": "FileNotFoundError",
     "evalue": "[Errno 2] No such file or directory: 'dataset/ccm/entity.txt'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mFileNotFoundError\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-2-67cd508612d7>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     22\u001b[0m \u001b[0mvalid_vocab_start\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;36m0\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     23\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 24\u001b[0;31m \u001b[0;32mwith\u001b[0m \u001b[0mopen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtarget_folder\u001b[0m\u001b[0;34m+\u001b[0m\u001b[0;34m'entity.txt'\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m'w+'\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mencoding\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m'utf-8'\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0mfout\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     25\u001b[0m     \u001b[0;32mwith\u001b[0m \u001b[0mopen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mentity_vocab_path\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m'r+'\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mencoding\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m'utf-8'\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0mfin\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     26\u001b[0m         \u001b[0mfout\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mwrite\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m''\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mjoin\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfin\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mreadlines\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mFileNotFoundError\u001b[0m: [Errno 2] No such file or directory: 'dataset/ccm/entity.txt'"
     ]
    }
   ],
   "source": [
    "# Your dataset folders\n",
    "source_folder = 'dataset/base/'\n",
    "target_folder = 'dataset/ccm/' \n",
    "\n",
    "# Knowlege Fact File, you can reuse the released 'naf_fact_vocab.txt'.\n",
    "# All facts are directional; hence if you want use other KBs, please follow the format:\n",
    "# Five items per line: 0:entity_in_post, 1:entity_in_response, 2 head_entity, 3 relation, 4 tail_entity,\n",
    "# where 0=2&1=3 or 0=3&1=2, entity_in_post/response is only used to match the post/response.\n",
    "# the first line should be '#N #N #NH #NF #NT'\n",
    "fact_vocab_path = 'concept/fact_vocab.txt'\n",
    "\n",
    "# Your entity/relation vocabs, \n",
    "# you can reuse the released 'entity/relation_vocab.txt' if you have removed all special tokens, i.e, #XXX.\n",
    "\n",
    "entity_vocab_path = 'concept/entity.txt'\n",
    "relation_vocab_path = 'concept/relation.txt'\n",
    "file_map = {\n",
    "      'train': 'trainset.txt',\n",
    "    'test': 'testset.txt',\n",
    "    'dev': 'validset.txt',\n",
    "}\n",
    "valid_vocab_start = 0\n",
    "\n",
    "with open(target_folder+'entity.txt','w+',encoding='utf-8') as fout:\n",
    "    with open(entity_vocab_path,'r+',encoding='utf-8') as fin:\n",
    "        fout.write(''.join(fin.readlines()))\n",
    "    \n",
    "\n",
    "with open(target_folder+'relation.txt','w+',encoding='utf-8') as fout:\n",
    "    with open(relation_vocab_path,'r+',encoding='utf-8') as fin:\n",
    "        fout.write(''.join(fin.readlines()))\n",
    "    \n",
    "    \n",
    "with open(target_folder+'fact_vocab.txt','w+',encoding='utf-8') as fout:\n",
    "    with open(fact_vocab_path,'r+',encoding='utf-8') as fin:\n",
    "        fout.write(''.join(fin.readlines()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 构建索引\n",
    "from collections import defaultdict\n",
    "\n",
    "head_fact_set = defaultdict(set)\n",
    "tail_fact_set = defaultdict(set)\n",
    "\n",
    "fact_vocabs = [x.strip('\\r\\n').split() for x in open(fact_vocab_path).readlines()]\n",
    "\n",
    "for fid, fact in enumerate(fact_vocabs):\n",
    "    head_entity = fact[0]\n",
    "    tail_entity = fact[1]\n",
    "    head_fact_set[head_entity].add(fid)\n",
    "    tail_fact_set[tail_entity].add(fid)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "entity_list = []\n",
    "entity_to_id = dict()\n",
    "with open(entity_vocab_path,'r',encoding='utf-8') as fin:\n",
    "    for line in fin.readlines():\n",
    "        items = line.split()\n",
    "        entity_list.append(items[0])\n",
    "        entity_to_id[items[0]] = len(entity_to_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def match_to_kb(post, response):\n",
    "    post_fid_set = set()\n",
    "    response_fid_set = set()\n",
    "    \n",
    "    for word in post:\n",
    "        for fid in head_fact_set[word]:\n",
    "            post_fid_set.add(fid)\n",
    "    \n",
    "    for word in response:\n",
    "        for fid in tail_fact_set[word]:\n",
    "            response_fid_set.add(fid)\n",
    "    \n",
    "    return post_fid_set, post_fid_set & response_fid_set\n",
    "\n",
    "def match_related(post):\n",
    "    post_fid_set = set()\n",
    "    \n",
    "    for word in post:\n",
    "        for fid in head_fact_set[word]:\n",
    "            post_fid_set.add(fid)\n",
    "    \n",
    "    return post_fid_set\n",
    "\n",
    "\n",
    "for dataset in ['train', 'dev', 'test']:\n",
    "    posts = [x.strip('\\r\\n').split() for x in open('%s%s.%s' % (source_folder, dataset, 'src')).readlines()]\n",
    "    responses = [x.strip('\\r\\n').split() for x in open('%s%s.%s' % (source_folder, dataset, 'tgt')).readlines()]\n",
    "    with open(target_folder+dataset+'.matched','w+',encoding='utf-8') as fmatched:\n",
    "        with open(target_folder+dataset+'.related','w+',encoding='utf-8') as frelated:\n",
    "            for i in range(len(posts)):\n",
    "                related_fact_ids, matched_fact_ids  = match_to_kb(posts[i], responses[i])\n",
    "                fmatched.write('%s\\n' % ' '.join([str(x) for x in list(matched_fact_ids)]))\n",
    "                frelated.write('%s\\n' % ' '.join([str(x) for x in list(related_fact_ids)]))\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "ename": "FileNotFoundError",
     "evalue": "[Errno 2] No such file or directory: 'dataset/base/train.src'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mFileNotFoundError\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-3-6c852c016eb5>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m    136\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    137\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mfile\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mfile_map\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mkeys\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 138\u001b[0;31m     \u001b[0;32mwith\u001b[0m \u001b[0mopen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0msource_folder\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mfile\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0;34m'.src'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m'r+'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mencoding\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m'utf-8'\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0mfin\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    139\u001b[0m         \u001b[0mposts\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstrip\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'\\r\\n'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msplit\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mx\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mfin\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mreadlines\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    140\u001b[0m \u001b[0;31m#         print(posts[0:10])\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mFileNotFoundError\u001b[0m: [Errno 2] No such file or directory: 'dataset/base/train.src'"
     ]
    }
   ],
   "source": [
    "from multiprocessing.pool import Pool\n",
    "import random\n",
    "import json\n",
    "from collections import defaultdict\n",
    "\n",
    "MAX_FACT = 50\n",
    "MAX_GOLDEN = 0\n",
    "\n",
    "def match_job(bulk_id, posts, responses, matched, related):\n",
    "    \n",
    "    print('job%d has started, 1st word: %s' % (bulk_id, posts[0][0]))\n",
    "    \n",
    "    random.seed(MAX_FACT)\n",
    "\n",
    "    cache = []\n",
    "    for post, response, matched, related in zip(posts, responses, matched, related):\n",
    "        data = dict()\n",
    "        \n",
    "        data['post'] = post\n",
    "        data['response'] = response\n",
    "        \n",
    "        all_facts = matched + related\n",
    "        selected_triplets = random.sample(all_facts, min(len(all_facts), MAX_FACT))\n",
    "        \n",
    "        matched_set = set(matched)\n",
    "        related_set = set(related)\n",
    "        \n",
    "        matched_triplets = []\n",
    "        related_triplets = []\n",
    "        \n",
    "        for triplet in selected_triplets:\n",
    "            if triplet in matched_set:\n",
    "                matched_triplets.append(triplet)\n",
    "            else:\n",
    "                related_triplets.append(triplet)\n",
    "                \n",
    "        if len(matched_triplets) == 0 and len(matched_set) > 0:\n",
    "            related_triplets = related_triplets[1:]\n",
    "            matched_triplets = random.sample(matched, 1)\n",
    "                \n",
    "        data['match_triples'] = matched_triplets\n",
    "            \n",
    "        used_triplets = matched_triplets + related_triplets\n",
    "        assert len(used_triplets) <= MAX_FACT \n",
    "        \n",
    "    \n",
    "        head_index = defaultdict(list)\n",
    "        tail_index = defaultdict(list)\n",
    "        for fid in used_triplets:\n",
    "            fact =  fact_vocabs[fid]\n",
    "            head_entity = fact[0]\n",
    "            tail_entity = fact[1]\n",
    "            head_index[head_entity].append(fid)\n",
    "            tail_index[tail_entity].append(fid)\n",
    "        matched_triplets = set(matched_triplets)\n",
    "        matched_triplets_dict = defaultdict(list)\n",
    "        matched_index_map = defaultdict(list)\n",
    "        \n",
    "        post_triples = []\n",
    "        all_triples = []\n",
    "        all_entities = []\n",
    "        \n",
    "        all_golden_triples = []\n",
    "        all_golden_entities = []\n",
    "        count = 0\n",
    "        for word in post:\n",
    "            if word in head_index:\n",
    "                count += 1\n",
    "                post_triples.append(count)\n",
    "                tmp_triplets = []\n",
    "                tmp_entities = []\n",
    "                golden_triplets = []\n",
    "                golden_entities = []\n",
    "                for fid in head_index[word]:\n",
    "                    if fid in matched_triplets :\n",
    "                        item  = fact_vocabs[fid]\n",
    "                        source_entity = item[0]\n",
    "                        target_entity = item[1]\n",
    "                        assert word == source_entity\n",
    "                        matched_index_map[fid].append((count, len(golden_triplets)))\n",
    "                        matched_triplets_dict[target_entity].append(fid)\n",
    "                        golden_triplets.append(fid)\n",
    "                        golden_entities.append(entity_to_id[item[1]])\n",
    "                    else:\n",
    "                        item  = fact_vocabs[fid]\n",
    "                        tmp_triplets.append(fid)\n",
    "                        tmp_entities.append(entity_to_id[item[1]])\n",
    "                \n",
    "                        \n",
    "                    \n",
    "                all_triples.append(tmp_triplets)\n",
    "                all_entities.append(tmp_entities)\n",
    "        \n",
    "                all_golden_triples.append(golden_triplets)\n",
    "                all_golden_entities.append(golden_entities)\n",
    "                assert len(tmp_triplets) == len(set(tmp_triplets))\n",
    "            else:\n",
    "                post_triples.append(0)\n",
    "                \n",
    "        for i in range(len(all_triples)):\n",
    "            all_triples[i] = all_golden_triples[i] +  all_triples[i]\n",
    "            all_entities[i] = all_golden_entities[i] +  all_entities[i]\n",
    "        \n",
    "        data['post_triples'] = post_triples\n",
    "        data['all_triples'] =  all_triples\n",
    "        data['all_entities'] =  all_entities\n",
    "        \n",
    "        # Response Triplets\n",
    "        response_triplets = []\n",
    "        match_index = []\n",
    "        for word in response:\n",
    "            if word in matched_triplets_dict:\n",
    "                matched_triplet = random.sample(matched_triplets_dict[word],1)[0]\n",
    "                response_triplets.append(matched_triplet)\n",
    "                if len(matched_index_map[matched_triplet]) > 0: # 可能有被抛弃的Golden Entity\n",
    "                    tmp = random.sample(matched_index_map[matched_triplet],1)[0]\n",
    "                    match_index.append([tmp[0],tmp[1]])\n",
    "                    assert entity_list[all_entities[tmp[0]-1][tmp[1]]] == word, '%s %s' % (entity_list[all_entities[tmp[0]-1][tmp[1]]], word)\n",
    "                else:\n",
    "                    response_triplets.append(-1)\n",
    "                    match_index.append([-1,-1])\n",
    "            else:\n",
    "                response_triplets.append(-1)\n",
    "                match_index.append([-1,-1])\n",
    "        data['response_triples'] = response_triplets\n",
    "        data['match_index'] = match_index\n",
    "        cache.append(data)\n",
    "        \n",
    "        assert len(data['response_triples'])  == len(data['match_index'])\n",
    "        assert len(data['response'])  == len(data['match_index'])\n",
    "#     print(' Cache Count and Line Number %d %d' % (len(cache), len(posts)))\n",
    "    assert len(posts) == len(cache) , '%d %d' % (len(cache), len(posts))\n",
    "    print('job%d has done, 1st word: %s' % (bulk_id, posts[0][0]))\n",
    "    return cache\n",
    "\n",
    "\n",
    "for file in file_map.keys():\n",
    "    with open(source_folder + file + '.src', 'r+', encoding='utf-8') as fin:\n",
    "        posts = [x.strip('\\r\\n').split() for x in fin.readlines()]\n",
    "#         print(posts[0:10])\n",
    "    print('Loaded-1')\n",
    "    with open(source_folder + file + '.tgt', 'r+', encoding='utf-8') as fin:\n",
    "        responses = [x.strip('\\r\\n').split() for x in fin.readlines()]\n",
    "#         print(responses[0:10])\n",
    "    print('Loaded-2')\n",
    "    with open(target_folder + file + '.matched', 'r+', encoding='utf-8') as fin:\n",
    "        matched = [x.strip('\\r\\n').split() for x in fin.readlines()]\n",
    "        matched = [[int(x) for x in y] for y in matched]\n",
    "#          print(matched[0:10])\n",
    "    print('Loaded-3')\n",
    "    with open(target_folder + file + '.related', 'r+', encoding='utf-8') as fin:\n",
    "        related = [x.strip('\\r\\n').split() for x in fin.readlines()]\n",
    "        related = [[int(x) for x in y] for y in related]\n",
    "#          print(related[0:10])\n",
    "    print('Loaded-4')\n",
    "    \n",
    "    line_number = len(posts)\n",
    "    num_thread = 3\n",
    "    bulk_num = line_number // num_thread + 1\n",
    "    pool = Pool(num_thread)\n",
    "    jobs = []\n",
    "    cache = []\n",
    "    for i in range(0, line_number, bulk_num):\n",
    "#         match_job(i, posts[i:i+bulk_num],responses[i:i+bulk_num],matched[i:i+bulk_num],related[i:i+bulk_num])\n",
    "        job = pool.apply_async(match_job, (i, posts[i:i+bulk_num],responses[i:i+bulk_num],matched[i:i+bulk_num],related[i:i+bulk_num]))\n",
    "        jobs.append(job)\n",
    "    pool.close()\n",
    "    pool.join()\n",
    "    \n",
    "    del posts\n",
    "    del responses\n",
    "    del matched\n",
    "    del related\n",
    "    \n",
    "    with open(target_folder+file_map[file], 'w+', encoding='utf-8') as fout:\n",
    "        for job in jobs:\n",
    "            for data in job.get():\n",
    "                line = json.dumps(data, ensure_ascii=False)\n",
    "                fout.write(line+'\\n')\n",
    "    print('Done',file)\n",
    "\n",
    "\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "ename": "FileNotFoundError",
     "evalue": "[Errno 2] No such file or directory: 'concept/fact_vocab.txt'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mFileNotFoundError\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-4-1c4b8b3e47f6>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      4\u001b[0m \u001b[0mkb_dict\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mdefaultdict\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mlist\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 6\u001b[0;31m \u001b[0;32mwith\u001b[0m \u001b[0mopen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfact_vocab_path\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m'r'\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mencoding\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m'utf-8'\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0mfin\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      7\u001b[0m     \u001b[0;32mfor\u001b[0m \u001b[0mi\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mline\u001b[0m \u001b[0;32min\u001b[0m \u001b[0menumerate\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfin\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mreadlines\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      8\u001b[0m         \u001b[0mitems\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mline\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstrip\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'\\r\\n'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msplit\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m' '\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mFileNotFoundError\u001b[0m: [Errno 2] No such file or directory: 'concept/fact_vocab.txt'"
     ]
    }
   ],
   "source": [
    "\n",
    "csk_triples = []\n",
    "csk_entities = []\n",
    "raw_vocab = dict()\n",
    "kb_dict = defaultdict(list)\n",
    "\n",
    "with open(fact_vocab_path,'r',encoding='utf-8') as fin:\n",
    "    for i,line in enumerate(fin.readlines()):\n",
    "        items = line.strip('\\r\\n').split(' ')\n",
    "        fact = '%s, %s, %s' % (items[2], items[3], items[4])\n",
    "        csk_triples.append(fact)\n",
    "        kb_dict[items[0]].append(fact)\n",
    "        kb_dict[items[2]].append(fact)\n",
    "        \n",
    "print(csk_triples[0:10])\n",
    "\n",
    "with open(entity_vocab_path,'r',encoding='utf-8') as fin:\n",
    "        for line in fin.readlines():\n",
    "            items = line.split()\n",
    "            csk_entities.append(items[0])\n",
    "print(csk_entities[0:10])\n",
    "\n",
    "with open(source_folder+'vocab.txt', 'r+') as fin:\n",
    "    vocabs = [x.strip('\\r\\n') for x in fin.readlines()]\n",
    "    for item in vocabs[valid_vocab_start:]:\n",
    "        raw_vocab[item] = 9999\n",
    "    print(len(vocabs))\n",
    "print(kb_dict['0分'])\n",
    "\n",
    "data = {\n",
    "    'csk_triples':csk_triples,\n",
    "    'csk_entities':csk_entities,\n",
    "    'vocab_dict':raw_vocab,\n",
    "    'dict_csk':kb_dict,\n",
    "}\n",
    "\n",
    "json.dump(data, open(target_folder+'resource.txt', 'w+', encoding='utf-8'), ensure_ascii=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
